import os
import pymysql
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

def reset_database():
    # 连接到MySQL（不指定数据库）
    conn = pymysql.connect(
        host=os.getenv("DB_HOST", "localhost"),
        port=int(os.getenv("DB_PORT", "3307")),
        user=os.getenv("DB_USER", "root"),
        password=os.getenv("DB_PASSWORD", "123456"),
        charset='utf8mb4'
    )
    
    try:
        with conn.cursor() as cursor:
            # 删除数据库（如果存在）
            cursor.execute("DROP DATABASE IF EXISTS chat_app")
            print("✅ 已删除旧数据库")
            
            # 读取并执行初始化SQL文件
            with open('init_db.sql', 'r', encoding='utf-8') as f:
                sql_commands = f.read()
            
            # 分割并执行多个SQL命令
            for command in sql_commands.split(';'):
                if command.strip():
                    cursor.execute(command)
                    conn.commit()
            
            print("✅ 已成功重建数据库")
            
    except Exception as e:
        print(f"❌ 重置数据库时出错：{str(e)}")
        raise
    finally:
        conn.close()

if __name__ == '__main__':
    print("开始重置数据库...")
    reset_database()
    print("数据库重置完成！") 